#include "DriverCallbacks.h"
#include "Defs.h"
#include "IpcDefs.h"
#include <wdm.h>
#include <Strsafe.h>

#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(PAGE, DriverUnload)

//injects device before the \\Device\\KeyboardClass0 class device driver
NTSTATUS InstallFilter( __in const PDRIVER_OBJECT pDriverObject)
{
    PDEVICE_OBJECT pKeyboardDevice;
    CHECK_AND_RETURN(IoCreateDevice(pDriverObject, sizeof(DeviceExtension), NULL, FILE_DEVICE_KEYBOARD, 0, FALSE, &pKeyboardDevice));
    pKeyboardDevice->Flags = pKeyboardDevice->Flags | (DO_BUFFERED_IO | DO_POWER_PAGABLE);
    pKeyboardDevice->Flags = pKeyboardDevice->Flags & ~DO_DEVICE_INITIALIZING;

    RtlSecureZeroMemory(pKeyboardDevice->DeviceExtension, sizeof(DeviceExtension));
    DeviceExtension* pdx = static_cast<DeviceExtension*>(pKeyboardDevice->DeviceExtension);
    UNICODE_STRING fidoName;
    RtlInitUnicodeString(&fidoName, L"\\Device\\KeyboardClass0");
    NTSTATUS status = IoAttachDevice(pKeyboardDevice, &fidoName, &pdx->nextDeviceObject);
    if(!NT_SUCCESS(status))
    {
        IoDeleteDevice(pKeyboardDevice);
    }
    return status;
}

NTSTATUS DriverEntry(__in PDRIVER_OBJECT  pDriverObject, __in PUNICODE_STRING  registryPath)
{   
    for (unsigned int i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; ++i)
    {
        pDriverObject->MajorFunction[i] = DispatchThruIO;        
    }

    pDriverObject->DriverUnload               = DriverUnload;
    pDriverObject->MajorFunction[IRP_MJ_READ] = DispatchRead;

    if(!IoIsWdmVersionAvailable(6, 0))
    {
        pDriverObject->MajorFunction[IRP_MJ_POWER] = DispatchThruPO;
    }

    return STATUS_SUCCESS;
}

NTSTATUS AddDevice(IN PDRIVER_OBJECT pDriverObject, IN PDEVICE_OBJECT PhysicalDeviceObject)
{
    CHECK_AND_RETURN(InstallFilter(pDriverObject));

    WCHAR namebuf[64];
    static LONG numextra = 0;

    StringCchPrintfW(namebuf, sizeof namebuf, L"%s%d", ipc_common::g_edoFullPath, InterlockedIncrement(&numextra));
    UNICODE_STRING edoname;
    RtlInitUnicodeString(&edoname, namebuf);
    PDEVICE_OBJECT edo;
    IoCreateDevice(pDriverObject, sizeof(ExtraDeviceExtension), &edoname, FILE_DEVICE_UNKNOWN, 0, FALSE, &edo);
    ExtraDeviceExtension* edx = static_cast<ExtraDeviceExtension*>(edo->DeviceExtension);
    DeviceExtension* pdx = static_cast<DeviceExtension*>(pDriverObject->DeviceObject->DeviceExtension);
    edx->pdx = pdx;
    pdx->edx = edx;

    return InstallFilter(pDriverObject);
}